# Author: lukas stoetzer
# Project: Dynamic Multinomial Dirichlet Model
## Content:
##    * update_alpha updates alpha of dirichlet distribution
##    * rdirichlet draws samples from dirichlet distribution
##    * get_ci_pi Get CI of probabilities for alpha values
##    * est_d Estimates discount Factor using maximum Liklihood
##    * for_alpha Forecast Alpha to election day
##    * get_coalop Estimate Coalition Majorities from each time-point
##    * est_gamma Estimates Relationship between Election results to create prior for Forecast


# %%%%%%%%%%%%% --------------------- %%%%%%%%%%%%%%% #

  update_alpha <- function(
                  Y, # Matrix of Support ncol=nParties nrow=nTime
                  palpha, # Prior Alpha
                  d=0.99 # Discount Factor
                  ){

    # Set missing Y to 0
      Y[is.na(Y)] <- 0

    # Result matrix
      alpha <- matrix(NA, nrow=nrow(Y)+1, ncol=ncol(Y))
      alpha[1,] <- palpha

    # Update alpha
      for(t in 1:nrow(Y)){
        alpha[t+1,] <-Y[t,] + d*alpha[t,]
      }

    # Return Result
      return(alpha)

    }


# %%%%%%%%%%%%% --------------------- %%%%%%%%%%%%%%% #

# Dirichlet Stuff


  rdirichlet <- function(
                    N = 100,  # Numer of random draws
                    alpha    # alpha Paramter
  ){


    # One random draw from Dirichlet
      rd <- function(a) { y <- rgamma(length(a), a, 1); y / sum(y) }

    # Replicate N times
      draws <- t(replicate(N, rd(alpha), simplify = TRUE))

    # Return draws
      return(draws)

  }



  qdirichlet <- function(a, q, N = 1000) {

    # Take random draws
    tmp <- rdirichlet( N, a)

    # Calculate quantiles
    quan <- apply(tmp,2, function(i) quantile(i,q, na.rm=T))

    # return
    return(t(quan))

  }

  ddirichlet <- function(x,alpha, log = FALSE){

    # Checks
    if (missing(x))
      stop("x is a required argument.")
    if (missing(alpha))
      stop("alpha is a required argument.")
    if (!is.matrix(x))
      x <- rbind(x)
    if (!is.matrix(alpha))
      alpha <- matrix(alpha, nrow(x), length(alpha), byrow = TRUE)
    if (any(rowSums(x) != 1))
      x/rowSums(x)
    if (any(x < 0))
      stop("x must be non-negative.")


    # log-density
    dens <- as.vector(lgamma(rowSums(alpha)) - rowSums(lgamma(alpha)) +
                        rowSums((alpha - 1) * log(x)))
    # Unlog
    if (log == FALSE)  dens <- exp(dens)

    # return
    return(dens)


  }

# %%%%%%%%%%%%% --------------------- %%%%%%%%%%%%%%% #


  get_ci_pi <- function(
            alpha,  # Alpha
            Nrep=1000,    # alpha Paramter
            ci=c(0.05,0.95),
            pty_nam
  ){


    # Require
      require(reshape2)

    # calc Mean pi
      sp <- apply(alpha,1,sum)
      pi_mean <- apply(alpha,2,function(x) x/sp)
      colnames(pi_mean) <- as.character(paste("mean",1:ncol(pi_mean),sep="."))

    # cacl quantiles
      pi_ci <- t(apply(alpha,1,qdirichlet,q=ci, N = Nrep))
      colnames(pi_ci) <- as.character(apply(expand.grid(1:ncol(alpha),paste("ci",ci*100,sep="")),1,function(x) paste(rev(x),collapse = ".")))

    # Turn into Data-frame
      d <- as.data.frame(cbind(1:nrow(pi_mean),pi_mean,pi_ci))
      colnames(d)[1] <- "time"

    # Reshape
      dl <- reshape(d,direction="long",varying=2:ncol(d),sep=".")

    # Name Parties
      colnames(dl)[1] <- c("pty_nam")
      dl$pty_nam <- as.factor(dl$pty_nam)
      levels(dl$pty_nam) <- pty_nam

    # Rename id
      colnames(dl)[which(colnames(dl)=="id")] <- "time"

    # Return ci
      return(dl)

  }

  # %%%%%%%%%%%%% --------------------- %%%%%%%%%%%%%%% #

  # Multinomial-Dirichlet Stuff


  qdirichlet <- function(a, q, N = 1000) {

    # Take random draws
    tmp <- rdirichlet( N, a)

    # Calculate quantiles
    quan <- apply(tmp,2, function(i) quantile(i,q))

    # return
    return(t(quan))

  }


# %%%%%%%%%%%%% --------------------- %%%%%%%%%%%%%%% #


  est_d <- function(y, n, a0){

    # Likelihood for lik(Y|d)
    lik_yt_Dtm1 <- function(d,Y,N,a0){

      # Information
      T <- nrow(Y) # Number time points
      K <- ncol(Y) # Number of Parties
      mY <- apply(Y,1,function(y) any(is.na(y))) # missing Data

      # Update alpha given discount factor d
      alpha <- matrix(NA, nrow=T, ncol=K)


      if(!mY[1]){
        alpha[1,] <- Y[1,] + d*a0 # Data Point
      } else {
        alpha[1,] <- d*a0 # Missing Data
      }

      for(t in 2:T){
        if(!mY[t]){
          alpha[t,] <- Y[t,] + d*alpha[t-1,] # Data Point
        } else {
          alpha[t,] <- d*alpha[t-1,]# Missing Data
        }
      }

      # logliklihood time point t
      logpYt <- 0

      for(t in which(!mY)[-1]){
        logpYt <- logpYt +
          lfactorial(N[t]) + lgamma(sum(d*alpha[t-1,])) - lgamma(N[t] + sum(d*alpha[t-1,])) +
          sum(lgamma(Y[t,] + d*alpha[t-1,]) - lfactorial(N[t]) - lgamma(d*alpha[t-1,]))
      }

      # Return sum log
      return(logpYt)

    }

    # Optimise
    res <- optim(par = 0.5,lik_yt_Dtm1,
                    method = "Brent",
                    hessian = T,
                    control=list(fnscale=-1),
                    lower=0.1,upper=1,
                    Y = y, N = n, a0 = a0)

    # Return result
    return(res)

  }

# %%%%%%%%%%%%% --------------------- %%%%%%%%%%%%%%% #


  for_alpha <- function(alpha,d){

    T <- nrow(alpha)
    for_a <- t(sapply(1:T, function(t) d^(T-t)*alpha[t,]))

    # Return result
    return(for_a)

  }

# %%%%%%%%%%%%% --------------------- %%%%%%%%%%%%%%% #


  get_coalop <- function(alpha,cbn = list(c(1,2)), Nrep = 1000, hurdel=0.05){

    # Define a Majority function
      majority_function <- function(share,sel_parties,tau){
        share <- ifelse(share<tau,0,share) # tau% Hurdel
        share <- share/sum(share) # Standardize
        coal_parties <- share[sel_parties]
        if(any(is.na(share))) {
          return(0)
        }
        if (any(coal_parties==0)){
          return(0) # if one of the coalition partners is below Hurdel return 0
        } else {
          return(as.numeric(sum(coal_parties) > 0.5)) # above 50 % return 1, below 0
        }
      }

    # Function to caclulate majority propability
      prob_majority <- function(sel, sample=draws){
        if(any(is.na(sel))) {
          return(NA)
        } else {
            pm <- mean(apply(sample,1,majority_function,sel_parties=sel,tau = hurdel))
            return(pm)
        }
      }

    # Take draws for alpha
      draws <- rdirichlet(Nrep,alpha)

    # Calculate for each of cobinations
      dfcbn <- do.call(rbind,lapply(cbn,prob_majority))

    # Return
      return(dfcbn)

  }


# %%%%%%%%%%%%% --------------------- %%%%%%%%%%%%%%% #

est_gamma <- function(res){

  # Log-Liklihood Function
  lik<- function(z,m){
    logl <- 0
    for(t in 2:nrow(m)){
      logl <- logl + ddirichlet(m[t,], z*m[t-1,], log = T)
    }
    return(logl)

  }

  # Optimize
  maxobj <- optimise(lik,m = res,interval = c(0.001,2000), maximum = T)

  # return
  return(maxobj)

}


